package com.minimalist.basic.config.interceptor;

import com.github.benmanes.caffeine.cache.Cache;
import com.github.benmanes.caffeine.cache.Caffeine;
import com.minimalist.basic.config.bucket.*;
import com.minimalist.basic.utils.Sm4Util;
import io.github.bucket4j.Bandwidth;
import io.github.bucket4j.Bucket;
import io.github.bucket4j.ConsumptionProbe;
import io.micrometer.core.instrument.MeterRegistry;
import io.micrometer.core.instrument.Timer;
import jakarta.servlet.http.HttpServletRequest;
import jakarta.servlet.http.HttpServletResponse;
import lombok.extern.slf4j.Slf4j;
import org.jetbrains.annotations.NotNull;
import org.springframework.http.HttpStatus;
import org.springframework.http.MediaType;
import org.springframework.scheduling.annotation.Scheduled;
import org.springframework.stereotype.Component;
import org.springframework.web.servlet.HandlerInterceptor;

import java.io.IOException;
import java.time.Duration;
import java.util.Map;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicReference;

/**
 * @author 11987
 */
@Slf4j
@Component
public class DynamicRateLimitInterceptor implements HandlerInterceptor, RateLimitConfigProvider {
    private final Cache<String, Bucket> bucketCache;
    private final SystemMetricsCollector metricsCollector;
    private final DynamicRateLimitCalculator calculator;
    private final AtomicReference<RateLimitConfig> currentConfig;
    private final RateLimitMetrics metrics;

    public DynamicRateLimitInterceptor(SystemMetricsCollector metricsCollector
            , DynamicRateLimitCalculator calculator, MeterRegistry meterRegistry) {
        this.metricsCollector = metricsCollector;
        this.calculator = calculator;
        this.currentConfig = new AtomicReference<>(
                new RateLimitConfig(100, Duration.ofMinutes(1))
        );
        this.bucketCache = Caffeine.newBuilder().expireAfterWrite(1, TimeUnit.HOURS)
                .build();
        this.metrics = new RateLimitMetrics(meterRegistry, this);
    }

    @Scheduled(fixedDelay = 10, timeUnit = TimeUnit.SECONDS)
    private void updateRateLimitConfig() {
        try {
            SystemMetrics metrics = metricsCollector.collectMetrics();
            RateLimitConfig newConfig = calculator.calculateLimit(metrics);

            RateLimitConfig oldConfig = currentConfig.get();
            if (hasSignificantChange(oldConfig, newConfig)) {
                currentConfig.set(newConfig);
                String logMessage = String.format("系统当前负载为: CPU负载=%.2f%% 内存占用=%.2f%% 速率限制更新为: %d/%ds",
                        metrics.getCpuLoad() * 100, metrics.getMemoryUsage() * 100,
                        newConfig.getLimit(), newConfig.getRefillDuration().getSeconds());
                log.info(logMessage);
                bucketCache.invalidateAll();
            }
        } catch (Exception e) {
            log.error("更新速率限制配置失败", e);
        }
    }

    private boolean hasSignificantChange(
            RateLimitConfig oldConfig, RateLimitConfig newConfig) {
        double limitChange = Math.abs(1.0 -
                (double) newConfig.getLimit() / oldConfig.getLimit());
        return limitChange > 0.2;
    }

    @Override
    public RateLimitConfig getRateLimitConfig() {
        return this.currentConfig.get();
    }

    @Override
    public boolean preHandle(HttpServletRequest request
            , @NotNull HttpServletResponse response
            , @NotNull Object handler) throws Exception {
        String path = request.getRequestURI();
        String method = request.getMethod();

        Timer.Sample timerSample = metrics.startTimer();
        boolean rateLimited = false;

        try {
            metrics.recordRequest();
            Map<String, Number> currentMetrics = metrics.getCurrentMetrics();
            log.info("当前速率限制值：{}，处理的请求总数：{}，超过速率限制的事件数：{}"
                    , currentMetrics.get("currentLimit"), currentMetrics.get("totalRequests")
                    , currentMetrics.get("rateLimitExceeded"));

            String clientId = getClientIdentifier(request);
            Bucket bucket = bucketCache.get(clientId, this::createBucket);

            ConsumptionProbe probe = bucket.tryConsumeAndReturnRemaining(1);

            if (probe.isConsumed()) {
                addRateLimitHeaders(response, probe);
                return true;
            }

            metrics.incrementRateLimitExceeded();
            handleRateLimitExceeded(response, probe);
            return false;
        } finally {
            metrics.stopTimer(timerSample, path, method, rateLimited);
        }
    }

    private Bucket createBucket(String clientId) {
        RateLimitConfig config = currentConfig.get();
        Bandwidth build = Bandwidth.builder().capacity(config.getLimit())
                .refillIntervally(config.getLimit(), config.getRefillDuration()).build();

        return Bucket.builder().addLimit(build).build();
    }

    private String getClientIdentifier(HttpServletRequest request) {
        return request.getRemoteAddr();
    }

    private void addRateLimitHeaders(HttpServletResponse response, ConsumptionProbe probe) {
        RateLimitConfig config = currentConfig.get();
        response.addHeader("X-Rate-Limit-Limit",
                String.valueOf(config.getLimit()));
        response.addHeader("X-Rate-Limit-Remaining",
                String.valueOf(probe.getRemainingTokens()));
        response.addHeader("X-Rate-Limit-Reset",
                String.valueOf(probe.getNanosToWaitForRefill() /
                        1_000_000_000));
    }

    private void handleRateLimitExceeded(HttpServletResponse response
            , ConsumptionProbe probe) throws IOException {
        response.setStatus(HttpStatus.TOO_MANY_REQUESTS.value());
        response.setContentType(MediaType.ALL_VALUE);

        String errorMessage = String.format("请求过于频繁，请等待 %d 秒后重试！",
                probe.getNanosToWaitForRefill() / 1_000_000_000);
        String encrypt = Sm4Util.encrypt(errorMessage);
        response.getWriter().write(encrypt);
    }
}
